import os, sys
import logging
from pathlib import Path
sys.path.append(str(Path(__file__).parent.parent.parent))

import re
import json
from langgraph.pregel.main import Graph
from pydantic import BaseModel, Field
from typing import AsyncGenerator, Dict, Any, List
import asyncio

from langchain_core.tools import tool
from langchain_deepseek import ChatDeepSeek
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import create_react_agent
from langgraph.checkpoint.memory import InMemorySaver
from langchain_core.messages import HumanMessage, AIMessage, ToolMessage

import json
import pymysql
from pymysql.cursors import DictCursor

from insight_agent.mcp_server_statistics.src.agent_prompts import agent_templates
from insight_agent.utils.http_post import update_insight_result, inform_insight_thinking

from insight_agent.mcp_server_statistics.config import DS_CONFIG, QWEN_CONFIG, DB_CONFIG, LLMConfig


# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler('agent_statistics.log', encoding='utf-8')
    ]
)
log = logging.getLogger(__name__)

current_dir = Path(__file__).parent





# 简单的事件队列，供流式输出使用
# 流式处理在stream_agent_response函数中处理

# 创建SQL查询工具
description = """
当用户需要进行数据库查询工作时，请调用该函数。
该函数用sqlite运行一段SQL代码，完成数据查询相关工作，
本函数只负责运行SQL代码并进行数据查询，若要进行数据提取，则使用另一个extract_data函数。
"""

# 初始化检查点存储，用于LangGraph状态管理
checkpointer = InMemorySaver()

# 数据库连接函数，避免全局连接
def get_db_connection():
    """获取数据库连接"""
    return pymysql.connect(
        host=DB_CONFIG.db_host,
        port=int(DB_CONFIG.db_port) if DB_CONFIG.db_port else 3306,
        user=DB_CONFIG.db_user,
        password=DB_CONFIG.db_password,
        database=DB_CONFIG.db_name,
        charset='utf8mb4',
        cursorclass=DictCursor
    )



# 定义结构化参数模型
class SQLQuerySchema(BaseModel):
    sql_query: str = Field(description=description)

# 封装为 LangGraph 工具
@tool(args_schema=SQLQuerySchema)
def sql_inquire(sql_query: str) -> str:
    """
    当用户需要进行数据库查询工作时，请调用该函数。
    该函数用MySQL运行一段SQL代码，完成数据查询相关工作，
    本函数只负责运行SQL代码并进行数据查询。
    :param sql_query: 字符串形式的SQL查询语句，用于执行对user_insight数据库中各张表进行查询，并获得各表中的各类相关信息
    :return：如果正确执行，是一个json字符串，字符串中包含sql_query、results、result_count三个key；
            如果执行失败，是一个json字符串，字符串包含sql_query、error、results、result_count四个key。
    """
    log.info("=" * 60)
    log.info(f"🔍 执行SQL查询: {sql_query}")
    
    connection = None
    try:
        # 获取新的数据库连接
        connection = get_db_connection()
        
        # 查询与检索数据
        cursor = connection.cursor()
        cursor.execute(sql_query)
        results = cursor.fetchall()

        # 预览前5条
        try:
            preview = results[:5]
        except Exception:
            preview = []

        # 将结果和SQL查询一起以特殊格式返回，便于后续解析
        response = {
            "sql_query": sql_query,
            "results": results,
            "result_count": len(results)
        }
        log.info(f"✅ 查询完成，返回 {len(results)} 条记录")
        return json.dumps(response, ensure_ascii=False)
    
    except Exception as e:
        error_response = {
            "sql_query": sql_query,
            "error": str(e),
            "results": [],
            "result_count": 0
        }
        log.error(f"❌ 查询失败: {str(e)}")
        return json.dumps(error_response, ensure_ascii=False)
    
    finally:
        # 确保连接被关闭
        log.info("=" * 60)
        if connection:
            connection.close()

# 创建工具列表
tools = [sql_inquire]

graph = None


def get_graph(config: LLMConfig) -> Graph:
    """创建LangGraph agent，确保不缓存结果"""
    prompt = agent_templates
    
    if not config.BASE_URL or not config.ALIYUN_API_KEY:
        raise ValueError("缺少LLM配置：请在环境变量中设置 BASE_URL 和 ALIYUN_API_KEY")
    
    # 创建模型实例，禁用缓存
    if config.MODEL.startswith('deepseek'):
        model = ChatDeepSeek(
            model=config.MODEL,
            base_url=config.BASE_URL,
            api_key=config.ALIYUN_API_KEY,
            api_base=config.BASE_URL,
            # 禁用缓存配置
            cache=False,
            # 添加随机性，避免缓存
            temperature=0.7,
            # 禁用流式缓存
            streaming=False
        )
    elif config.MODEL == 'qwen-max':
        model = ChatOpenAI(
            model="qwen-max",
            openai_api_base=config.BASE_URL,
            openai_api_key=config.ALIYUN_API_KEY,
            # 禁用缓存配置
            cache=False,
            # 添加随机性，避免缓存
            temperature=0.7,
            # 禁用流式缓存
            streaming=False
        )
    else:
        raise ValueError(f"不支持的模型: {config.MODEL}")
    
    # 创建LangGraph agent，使用检查点存储
    graph = create_react_agent(
        model=model, 
        tools=tools, 
        prompt=prompt,
        checkpointer=checkpointer,
        # max_iterations=2
    )
    
    return graph

from mcp.server.fastmcp import FastMCP

mcp = FastMCP("AgentStatisticsServer", port=5130)

# @mcp.tool()
async def statistics_info_oneway(user_input: str, env: str='dev'):
    """
    这是统算智能体（statistics agent）。
    :param user_input: 用户请求
    :return:

    【适用场景】
    电信领域存量客户的属性信息、指标-消费行为分析数据、用户所属标签分析的功能，例如：
    - “请统计北京市不同区中高端用户类型的平均ARPU值”
    - "请问中高端用户按照资费波动，被分为几类用户群，分别是什么，并且每个用户群的用户数量数是多少？",
    - "统计北京市不同区中高端用户类型的平均ARPU值",
    - "统计北京市不同区的银发用户数量，并统计其家宽在线时长的平均值",
    - "统计北京市不同区腰部用户的平均使用流量（MB）"

    【调用方式】
    1. 将用户的原始输入（包含数据指标、数据维度、对比意愿、分析需求等）作为 user_input 传递给本工具。
    2. 本工具会自动解析用户属性、标签、指标等元素，针对分析进行统计分析，并返回结构化的分析结果和简要总结。

    【返回内容】
    分析结果的 summary（简要结论）
    - 结构化的时序数据（如每日/每月的统计值、趋势线等）

    【注意事项】
    - 所有查询都不能涉及时间，有时间元素的本工具一概不处理。
    - 所有涉及地图的询问都不处理。
    """

    if env == 'test':
        inform_insight_thinking = lambda *args, **kwargs: None
        update_insight_result = lambda *args, **kwargs: None
    

    # 通过 HTTP 获取 main_agent 总结的与本智能体相关的内容
    # insight_request_response = get_insight_request("agent_statistics")  # 获取 insight_request 中 agent_statistics 相关的用户输入/请求
    # insight_payload = insight_request_response.get('insight_request', {})
    # user_request_summary = insight_payload.get('user_request_summary', '')
    # user_latest_request = insight_payload.get('user_latest_request') or user_request_summary
    
    log.info(f'💡传过来的用户输入:{user_input}')
    inform_insight_thinking(
        agent_type="agent_statistics", 
        thinking_step=f"用户说{user_input}"
    )
    # log.info(f'💡子智能体 agent_statistics 通过HTTP 收到用户请求：{user_latest_request}')
    # log.info(f'💡子智能体 agent_statistics 通过HTTP 收到用户请求总结：{user_request_summary}')

    question = user_input
    
    # log.info(f'💡子智能体 agent_statistics 根据用户输入生成了数据查询语句 {sql_query}')

    config = DS_CONFIG
    log.info(f"目前的DS_CONFIG信息:{config}")
    # try:
        # 每次创建新的graph实例，避免缓存
    agent = get_graph(config)
    
    # 生成唯一的线程ID，避免缓存
    import uuid
    thread_id = str(uuid.uuid4())
    
    config_dict = {
        "configurable": {
            "thread_id": thread_id
        }
    }

    log.info(f'你好你好你好')
    # 流式调用智能体
    ai_output = await agent.ainvoke(
        {"messages": [{"role": "user", "content": question}]},
        config_dict,
        stream_mode="values"
    )
    # log.info(f'💡子智能体 agent_statistics 返回的回答：{ai_output}')
        # 处理每个步骤
        #print(f"STEP: {step}\n\n")

    content = extract_json(ai_output['messages'][-1].content)
    inform_insight_thinking(
        agent_type="agent_statistics", 
        thinking_step=f"{content}"
    )
    # 给主智能体发送post请求
    update_insight_result(agent_type="agent_statistics", updated_result=content)
    return json.dumps(content, ensure_ascii=False)
    
@mcp.tool()
async def agent_statistics(question: str, env: str='dev'):
    """
    仅当用户输入提及“统计”时使用本工具，其余输入不使用本工具。
    :param user_input: 用户请求
    :return: 文字总结

    【适用场景】
    电信领域存量客户的属性信息、指标-消费行为分析数据、用户所属标签分析的功能，例如：
    - “请统计北京市不同区中高端用户类型的平均ARPU值”
    - "请问中高端用户按照资费波动，被分为几类用户群，分别是什么，并且每个用户群的用户数量数是多少？",
    - "统计北京市不同区的银发用户数量，并统计其家宽在线时长的平均值",
    - "统计北京市不同区腰部用户的平均使用流量（MB）"

    【调用方式】
    1. 将用户的原始输入（包含数据指标、数据维度、对比意愿、分析需求等）作为 user_input 传递给本工具。
    2. 本工具会自动解析用户属性、标签、指标等元素，针对分析进行统计分析，并返回结构化的分析结果和简要总结。

    【返回内容】
    - 分析结果的 summary（简要结论）

    """
    log.info(f"💡传过来的用户输入:{question}")
    log.info(f"💡传过来的环境:{env}")

    if env == 'dev':
        global inform_insight_thinking
        global update_insight_result
    elif env == 'test':
        inform_insight_thinking = lambda *args, **kwargs: print(kwargs['thinking_step'])
        update_insight_result = lambda *args, **kwargs: None

    # 生成唯一的线程ID，避免缓存
    import uuid
    thread_id = str(uuid.uuid4())
    
    config_dict = {
        "configurable": {
            "thread_id": thread_id
        }
    }

    try:
        # 每次创建新的graph实例，避免缓存
        log.info(f"💡DOGDOGDOG")
        try:
            agent = get_graph(DS_CONFIG)
            log.info(f"💡AGENT_INFO:{agent}")
        except Exception as e:
            log.error(f"💡AGENT_ERROR:{e}")

        log.info(f"💡agent:{agent}")
        

        try:
            inform_insight_thinking(
                agent_type="agent_statistics", 
                thinking_step=f"用户说{question}"
            )
        except Exception as e:
            log.error(f"💡INFORM_INSIGHT_THINKING_ERROR:{e}")
        # yield "🤔 正在分析您的问题..."
        
        # 流式调用智能体
        async for step in agent.astream(
            {"messages": [{"role": "user", "content": question}]},
            config_dict,
            stream_mode="values"
        ):
            # 处理每个步骤
            #print(f"STEP: {step}\n\n")
            if "messages" in step:
                message = step["messages"][-1]
                if isinstance(message, HumanMessage):
                    print(f"HumanMessage: {message.content}\n")
                elif isinstance(message, AIMessage):
                    if message.additional_kwargs.get('tool_calls', None):
                        last_tool_call = message.additional_kwargs.get('tool_calls')[-1]
                        tool_name = last_tool_call['function'].get('name', '')
                        tool_input = last_tool_call['function'].get('arguments', '')
                        str_num = min(40, len(tool_input)//2)
                        output_str = str(tool_input[:str_num]) + "..."
                        inform_insight_thinking(
                            agent_type="agent_statistics", 
                            thinking_step=f"我将使用 {tool_name}, 参数是 {output_str}\n"
                        )
                    elif message.content != '':
                        str_num = min(200000000, len(message.content))
                        output_str = str(message.content[:str_num]) + "..."
                        log.info(f"💡AIMessage: {output_str}")
                        inform_insight_thinking(
                            agent_type="agent_statistics", 
                            thinking_step=f"{output_str}\n"
                        )
                    else:
        
                        inform_insight_thinking(
                            agent_type="agent_statistics", 
                            thinking_step=f"AIMessage, 异类: {message}\n"
                        )
                        continue
                elif isinstance(message, ToolMessage):
                    str_num = min(40, len(message.content)//2)
                    output_str = str(message.content[:str_num]) + "..."
                    inform_insight_thinking(
                        agent_type="agent_statistics", 
                        thinking_step=f"工具调用 - {output_str}\n"
                    )
                else:
                    inform_insight_thinking(
                        agent_type="agent_statistics", 
                        thinking_step=f"🔍 Unknown message type: {type(message)}\n"
                    )
                continue

        ## 处理最终返回结果
        if step['messages'][-1].content != '':
            content = step['messages'][-1].content
            print("\n" * 5)
            print(f"💡即将进行json转换的json内容:\n\n {content}")
            json_content = extract_json(content)
            print(f"💡转换后的json类型:\n\n {type(json_content)}")
            try:
                json_dict = json.loads(json_content)
            except Exception as e:
                log.error(f"💡JSON转换失败: {e}")
                return json.dumps({"step_type": "error", "content": f"❌ 处理出错: {str(e)}"})
            else:
                update_insight_result(agent_type="agent_statistics", updated_result=json_dict)
                return json.dumps(json_dict, ensure_ascii=False)
            
            return json.dumps({"step_type": "error", "content": f"❌ 处理出错: {str(e)}"})
    except Exception as e:
        return json.dumps({"step_type": "error", "content": f"❌ 处理出错: {str(e)}"})
    

def extract_json(content: str):
    """从LLM响应中提取JSON"""
    # print("\n" * 2)
    # print(content)
    try:
        # 提取 ``` 之间的内容
        pattern = r'```(?:json)?\s*\n(.*?)\n```'
        match = re.search(pattern, content, re.DOTALL)

        print(f"💡match: {match}")
        
        if match:
            json_content = match.group(1).strip()
            log.info(f"💡提取的json内容:\n\n {json_content}")
            return json_content
        elif match is None:
            json_content = content
            try:
                json_dict = json.loads(json_content)
                return json_content
            except Exception as e:
                error_content = f"❌ JSON内容不match且JSON转换失败: {e}"
                return json.loads(f"{error_content: {error_content}}")
        
    except Exception as e:
        error_content = f"❌ JSON提取失败: {e}"
        return json.loads(f"{error_content: {error_content}}")

# 同步调用函数（用于测试）
def invoke_agent_sync(question: str, config: LLMConfig = DS_CONFIG):
    """同步调用智能体，每次创建新实例避免缓存"""
    try:
        # 每次创建新的graph实例
        agent = get_graph(config)
        
        # 生成唯一的线程ID
        import uuid
        thread_id = str(uuid.uuid4())
        
        config_dict = {
            "configurable": {
                "thread_id": thread_id
            }
        }
        
        # 调用智能体
        response = agent.invoke(
            {"messages": [{"role": "user", "content": question}]},
            config_dict
        )
        log.info(response)
        response.keys()
        log.info(type(response))
        
        content = response['messages'][-1].content
        content_json = extract_json(content)

        return json.dumps(content_json, ensure_ascii=False)
        
    except Exception as e:
        log.error(f"❌ 调用失败: {str(e)}")
        return None


if __name__ == "__main__":
    # from pprint import pprint
    
    # # 测试问题
    # test_question = "请统计北京市不同区中高端用户类型的平均ARPU值"
    # # test_question = "请统计北京不同区中高端用户数"
    
    log.info("🚀 开始测试LangGraph智能体（无缓存模式）")
    # log.info(f"📝 测试问题: {test_question}")
    log.info("-" * 50)
    
    # # 选择测试模式
    # test_mode = "sync"  # 可选: "sync", "async"
    
    # if test_mode == "sync":
    #     # 使用同步调用函数
    #     print("🔄 使用同步调用模式")
    #     response = invoke_agent_sync(test_question, DS_CONFIG)
        
    #     if response:
    #         print("✅ 调用成功!")
    #         print("📊 响应内容:")
    #         print(response)
    #     else:
    #         print("❌ 调用失败!")
    
    # elif test_mode == "async":
    #     # 使用异步流式调用函数
    #     print("🔄 使用异步流式调用模式")
        
    #     async def run_async_test():
    #         async for chunk in stream_agent_response(test_question, DS_CONFIG):
    #             step_type = chunk.get("step_type", "unknown")
    #             content = chunk.get("content", "")
    #             #log.info(chunk.content, end='|', flush=True)
    #             log.info(content)
    #     asyncio.run(run_async_test())
        
    #     # 运行异步函数
       
    mcp.run(transport="streamable-http")
    
    
    # async def test():
    #     result = await statistics_info("请统计北京市不同区中高端用户类型的平均ARPU值")
    #     print(f"测试结果:{result}")
    # asyncio.run(test())

    log.info("-" * 50)
    log.info("🎯 测试完成!")